#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sokol_audio.h>
#include <sokol_log.h>
#include <assert.h>
#include "audio.h"

#define AUDIO_SAMPLE_RATE 96000
#define BUFFER_SIZE 8192
#define REQUESTED_FRAME_COUNT 2048

#define MEMORY_ORDER memory_order_seq_cst

static void saudio_callback(float* buffer, int num_frames, int num_channels, void* user_data) {
    AudioDriver* data = (AudioDriver*) user_data;

    size_t read_ptr = data->queue_read;
    size_t write_ptr = data->queue_write;

    size_t available;
    if (write_ptr >= read_ptr) {
        available = write_ptr - read_ptr;
    } else {
        available = BUFFER_SIZE - read_ptr + write_ptr;
    }

    if (available < num_frames) {
        memset(buffer, 0, num_frames * num_channels * sizeof(float));
        return;
    }

    float fsample;
    for (int i = 0; i < num_frames; i++) {
        fsample = data->audio_queue[read_ptr];
        read_ptr = (read_ptr + 1) % BUFFER_SIZE;

        if (fsample < -1.0f) fsample = -1.0f;
        if (fsample > 1.0f) fsample = 1.0f;
        buffer[0] = fsample;

        // fill channels
        for (int c = 1; c < data->channel_count; c++) {
            buffer[c] = buffer[0];
        }

        buffer += data->channel_count;
    }

    data->queue_read = read_ptr;
}

int audio_init(AudioDriver* audio_struct) {
    saudio_setup(&(saudio_desc){
        .logger.func = slog_func,
        .stream_userdata_cb = saudio_callback,
        .user_data = audio_struct,
        .sample_rate = AUDIO_SAMPLE_RATE,
        .num_channels = 2,
        .buffer_frames = REQUESTED_FRAME_COUNT
    });

    if (!saudio_isvalid()) return 1;

    audio_struct->channel_count = saudio_channels();
    audio_struct->sample_rate = saudio_sample_rate();
    
    audio_struct->audio_queue = malloc(BUFFER_SIZE * sizeof(float));
    audio_struct->queue_write = 0;
    audio_struct->queue_read = 0;

#ifdef AUDIO_WRITE_TO_FILE
    {
        char buf[8];

        uint32_t bps = sizeof(int16_t);

        audio_struct->_test = fopen("test.wav", "wb");
        fputs("RIFF", audio_struct->_test);
        fwrite("\0\0\0\0", 1, 4, audio_struct->_test); // file size, filled in later
        fputs("WAVE", audio_struct->_test);
        fwrite("fmt ", 4, 1, audio_struct->_test);

        // pcm mode
        *((uint32_t*)buf) = (uint32_t) 16;
        fwrite(buf, 4, 1, audio_struct->_test);

        // format type
        *((uint16_t*)buf) = (uint16_t) 1;
        fwrite(buf, 2, 1, audio_struct->_test);

        // num channels
        *((uint16_t*)buf) = (uint16_t) 2;
        fwrite(buf, 2, 1, audio_struct->_test);

        // sample rate
        *((uint32_t*)buf) = audio_struct->sample_rate;
        fwrite(buf, 4, 1, audio_struct->_test);

        // data rate
        *((uint32_t*)buf) = audio_struct->sample_rate * 2 * bps; 
        fwrite(buf, 4, 1, audio_struct->_test);

        // data block size
        *((uint16_t*)buf) = 2 * bps;
        fwrite(buf, 2, 1, audio_struct->_test);

        // bits per sample
        *((uint16_t*)buf) = bps * 8;
        fwrite(buf, 2, 1, audio_struct->_test);

        // data chunk header
        fputs("data", audio_struct->_test);

        // data section size
        *((uint32_t*)buf) = 0;
        fwrite(buf, 4, 1, audio_struct->_test);

        audio_struct->_samples_written = 0;
    }
#endif

    return 0;
}

void audio_destroy(AudioDriver* audio_struct) {
    saudio_shutdown();
    free(audio_struct->audio_queue);

#ifdef AUDIO_WRITE_TO_FILE
    {
        uint32_t chunk_size = audio_struct->_samples_written * sizeof(int16_t);
        uint32_t file_size = 36 + chunk_size;

        int err;
        err = fseek(audio_struct->_test, 4, SEEK_SET);
        assert(err == 0);
        fwrite(&file_size, 4, 1, audio_struct->_test);
        err = fseek(audio_struct->_test, 40, SEEK_SET);
        assert(err == 0);
        fwrite(&chunk_size, 4, 1, audio_struct->_test);

        fclose(audio_struct->_test);
    }
#endif
}

int audio_queue(AudioDriver* audio_struct, float* samples, size_t buffer_size) {
    size_t read_ptr = audio_struct->queue_read;
    size_t write_ptr = audio_struct->queue_write;
    
    size_t available;
    if (read_ptr > write_ptr) {
        available = read_ptr - write_ptr - 1;
    } else {
        available = BUFFER_SIZE - write_ptr + read_ptr - 1;
    }

    if (available < buffer_size) return 1;

    for (int i = 0; i < buffer_size; i++) {
        audio_struct->audio_queue[write_ptr] = samples[i];
        write_ptr = (write_ptr + 1) % BUFFER_SIZE;
    }

    audio_struct->queue_write = write_ptr;

#ifdef AUDIO_WRITE_TO_FILE
    {
        float fsample;
        for (int i = 0; i < buffer_size; i++) {
            fsample = samples[i];
            if (fsample < -1.0f) fsample = -1.0f;
            if (fsample > 1.0f) fsample = 1.0f;

            int16_t a = (int16_t)(fsample * INT16_MAX);
            fwrite(&a, 2, 1, audio_struct->_test);
            fwrite(&a, 2, 1, audio_struct->_test);
        }

        audio_struct->_samples_written += buffer_size * 2;
    }
#endif

    return 0;
}

size_t audio_available(const AudioDriver* audio_struct) {
    size_t read_ptr = audio_struct->queue_read;
    size_t write_ptr = audio_struct->queue_write;
    
    size_t available;
    if (read_ptr > write_ptr) {
        available = read_ptr - write_ptr - 1;
    } else {
        available = BUFFER_SIZE - write_ptr + read_ptr - 1;
    }
    
    return available;
}